FROM {{ cuda.base_image }} as base

{% if "x86_64" in cuda %}
FROM base as base-amd64

ENV NVARCH x86_64
{% if "requires" in cuda.x86_64 %}
ENV NVIDIA_REQUIRE_CUDA "{{ cuda.x86_64.requires }}"
{% endif %}
ENV NV_CUDA_CUDART_VERSION {{ cuda.x86_64.components.cudart.version }}

{% if cuda.x86_64.use_ml_repo %}
ENV NV_ML_REPO_ENABLED 1
ENV NV_ML_REPO_URL {{ cuda.ml_repo_url }}/${NVARCH}

{% endif -%}

{% if cuda.version.major | int == 8 %}
ENV NV_CUDA_NVRTC_VERSION {{ cuda.x86_64.components.nvrtc.version }}
ENV NV_CUDA_NVRTC_PACKAGE cuda-nvrtc-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_NVRTC_VERSION}
ENV NV_CUDA_NVGRAPH_VERSION {{ cuda.x86_64.components.nvgraph.version }}
ENV NV_CUDA_NVGRAPH_PACKAGE cuda-nvgraph-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_NVGRAPH_VERSION}
ENV NV_LIBCUSOLVER_VERSION {{ cuda.x86_64.components.libcusolver.version }}
ENV NV_LIBCUSOLVER_PACKAGE cuda-cusolver-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUSOLVER_VERSION}
ENV NV_LIBCUBLAS_VERSION {{ cuda.x86_64.components.libcublas.version }}
ENV NV_LIBCUBLAS_PACKAGE cuda-cublas-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUBLAS_VERSION}
ENV NV_LIBCUFFT_VERSION {{ cuda.x86_64.components.libcufft.version }}
ENV NV_LIBCUFFT_PACKAGE cuda-cufft-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUFFT_VERSION}
ENV NV_LIBCURAND_VERSION {{ cuda.x86_64.components.libcurand.version }}
ENV NV_LIBCURAND_PACKAGE cuda-curand-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCURAND_VERSION}
ENV NV_LIBCUSPARSE_VERSION {{ cuda.x86_64.components.libcusparse.version }}
ENV NV_LIBCUSPARSE_PACKAGE cuda-cusparse-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUSPARSE_VERSION}
ENV NV_LIBNPP_VERSION {{ cuda.x86_64.components.libnpp.version }}
ENV NV_LIBNPP_PACKAGE cuda-npp-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBNPP_VERSION}
{% endif -%}

{% endif %}
{% if "ppc64le" in cuda %}
FROM base as base-ppc64le

ENV NVARCH ppc64el
{% if "requires" in cuda.ppc64le %}
ENV NVIDIA_REQUIRE_CUDA "{{ cuda.ppc64le.requires }}"
{% endif %}
ENV NV_CUDA_CUDART_VERSION {{ cuda.ppc64le.components.cudart.version }}

{% if cuda.ppc64le.use_ml_repo %}
ENV NV_ML_REPO_ENABLED 1
ENV NV_ML_REPO_URL {{ cuda.ml_repo_url }}/${NVARCH}

{% endif %}
{% if cuda.version.major | int == 8 %}
ENV NV_CUDA_NVRTC_VERSION {{ cuda.ppc64le.components.nvrtc.version }}
ENV NV_CUDA_NVRTC_PACKAGE cuda-nvrtc-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_NVRTC_VERSION}
ENV NV_CUDA_NVGRAPH_VERSION {{ cuda.ppc64le.components.nvgraph.version }}
ENV NV_CUDA_NVGRAPH_PACKAGE cuda-nvgraph-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_NVGRAPH_VERSION}
ENV NV_LIBCUSOLVER_VERSION {{ cuda.ppc64le.components.libcusolver.version }}
ENV NV_LIBCUSOLVER_PACKAGE cuda-cusolver-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUSOLVER_VERSION}
ENV NV_LIBCUBLAS_VERSION {{ cuda.ppc64le.components.libcublas.version }}
ENV NV_LIBCUBLAS_PACKAGE cuda-cublas-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUBLAS_VERSION}
ENV NV_LIBCUFFT_VERSION {{ cuda.ppc64le.components.libcufft.version }}
ENV NV_LIBCUFFT_PACKAGE cuda-cufft-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUFFT_VERSION}
ENV NV_LIBCURAND_VERSION {{ cuda.ppc64le.components.libcurand.version }}
ENV NV_LIBCURAND_PACKAGE cuda-curand-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCURAND_VERSION}
ENV NV_LIBCUSPARSE_VERSION {{ cuda.ppc64le.components.libcusparse.version }}
ENV NV_LIBCUSPARSE_PACKAGE cuda-cusparse-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBCUSPARSE_VERSION}
ENV NV_LIBNPP_VERSION {{ cuda.ppc64le.components.libnpp.version }}
ENV NV_LIBNPP_PACKAGE cuda-npp-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_LIBNPP_VERSION}
{% endif -%}

{% endif %}
FROM base-${TARGETARCH}

ARG TARGETARCH

LABEL maintainer "NVIDIA CORPORATION <cudatools@nvidia.com>"

RUN apt-get update && apt-get install -y --no-install-recommends \
{% if "18.04" in cuda.os.version %}
    gnupg2 curl ca-certificates && \
    curl -fsSL {{ cuda.repo_url }}/${NVARCH}/7fa2af80.pub | apt-key add - && \
    echo "deb {{ cuda.repo_url }}/${NVARCH} /" > /etc/apt/sources.list.d/cuda.list && \
    if [ ! -z ${NV_ML_REPO_ENABLED} ]; then echo "deb ${NV_ML_REPO_URL} /" > /etc/apt/sources.list.d/nvidia-ml.list; fi && \
    apt-get purge --autoremove -y curl \
{% else %}
    ca-certificates apt-transport-https gnupg-curl && \
    NVIDIA_GPGKEY_SUM=d1be581509378368edeec8c1eb2958702feedf3bc3d17011adbf24efacce4ab5 && \
    NVIDIA_GPGKEY_FPR=ae09fe4bbd223a84b2ccfce3f60f4b3d7fa2af80 && \
    apt-key adv --fetch-keys {{ cuda.repo_url }}/${NVARCH}/7fa2af80.pub && \
    apt-key adv --export --no-emit-version -a $NVIDIA_GPGKEY_FPR | tail -n +5 > cudasign.pub && \
    echo "$NVIDIA_GPGKEY_SUM  cudasign.pub" | sha256sum -c --strict - && rm cudasign.pub && \
    echo "deb {{ cuda.repo_url }}/${NVARCH} /" > /etc/apt/sources.list.d/cuda.list && \
    if [ ! -z ${NV_ML_REPO_ENABLED} ]; then echo "deb ${NV_ML_REPO_URL} /" > /etc/apt/sources.list.d/nvidia-ml.list; fi && \
    apt-get purge --auto-remove -y gnupg-curl \
{% endif %}
    && rm -rf /var/lib/apt/lists/*

{% if "release_label" in cuda.version and cuda.version.release_label.split('.') | length > 2 %}
ENV CUDA_VERSION {{ cuda.version.release_label }}
{% else %}
ENV CUDA_VERSION {{ cuda.version.release_label }}.{{ cuda.version.build_version }}
{% endif %}

# For libraries in the cuda-compat-* package: https://docs.nvidia.com/cuda/eula/index.html#attachment-a
RUN apt-get update && apt-get install -y --no-install-recommends \
    cuda-cudart-{{ cuda.version.major }}-{{ cuda.version.minor }}=${NV_CUDA_CUDART_VERSION} \
{% if cuda.version.major | int >= 10 and cuda.arch != "arm64" %}
    cuda-compat-{{ cuda.version.major }}-{{ cuda.version.minor }} \
{% elif cuda.version.major | int == 8 %}
    ${NV_CUDA_NVRTC_PACKAGE} \
    ${NV_CUDA_NVGRAPH_PACKAGE} \
    ${NV_CUSOLVER_PACKAGE} \
    ${NV_LIBCUBLAS_PACKAGE} \
    ${NV_LIBCUFFT_PACKAGE} \
    ${NV_LIBCURAND_PACKAGE} \
    ${NV_LIBCUSPARSE_PACKAGE} \
    ${NV_LIBNPP_PACKAGE} \
{% endif %}
    && ln -s cuda-{{ cuda.version.major }}.{{ cuda.version.minor }} /usr/local/cuda && \
    rm -rf /var/lib/apt/lists/*

# Required for nvidia-docker v1
{% if cuda.version.major | int < 10 %}
LABEL com.nvidia.volumes.needed="nvidia_driver"
LABEL com.nvidia.cuda.version="${CUDA_VERSION}"
{% endif -%}
RUN echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf && \
    echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf

ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:${PATH}
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64

COPY NGC-DL-CONTAINER-LICENSE /

# nvidia-container-runtime
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
